# -*- coding: utf-8 -*-
"""
@Time ： 2022/9/19 15:19
@Auth ： GaoShuai
@File ：EvaluationMetrics.py
@IDE ：PyCharm
"""
import numpy as np
from skimage import measure
import torch


def computeSSIM(fake, real):
    if isinstance(fake, torch.Tensor):
        fake = fake.clone().detach().cpu().numpy().squeeze()
    if isinstance(real, torch.Tensor):
        real = real.clone().detach().cpu().numpy().squeeze()

    return measure.compare_ssim(fake, real)


def computePSNR(fake, real):
    if isinstance(fake, torch.Tensor):
        fake = fake.clone().detach().cpu().numpy().squeeze()
    if isinstance(real, torch.Tensor):
        real = real.clone().detach().cpu().numpy().squeeze()
    x, y = np.where(real != 0)  # Exclude background
    # x, y = np.where(real != 0)
    # x, y = np.where(real != np.min(fake))
    mse = np.mean(((fake[x][y] + 1) / 2. - (real[x][y] + 1) / 2.) ** 2)
    if mse < 1.0e-10:
        return 100
    else:
        PIXEL_MAX = 1
        return 20 * np.log10(PIXEL_MAX / np.sqrt(mse))


def computeMAE(fake, real):
    if isinstance(fake, torch.Tensor):
        fake = fake.clone().detach().cpu().numpy().squeeze()
    if isinstance(real, torch.Tensor):
        real = real.clone().detach().cpu().numpy().squeeze()
    x, y = np.where(real != 0)  # Exclude background
    # x, y = np.where(real != np.min(fake))
    # x, y = np.where(real != 0)
    mae = np.abs(fake[x, y] - real[x, y]).mean()
    return mae / 2  # from (-1,1) normaliz  to (0,1)
